package cn.da.shuai.cool.ai.search.client.util;

import cn.da.shuai.cool.ai.search.client.common.enums.SSEMsgType;
import cn.hutool.core.collection.CollectionUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class SSEServerUtil {

    private static final Map<String, SseEmitter> sseClients = new ConcurrentHashMap<>();

    public static SseEmitter connect(String userId) {
        SseEmitter sseEmitter = new SseEmitter(0L);
        sseEmitter.onTimeout(() -> {
            remove(userId);
            log.info("sse timeout: {}", userId);
        });
        sseEmitter.onError((e) -> {
            remove(userId);
            log.error("sse error: {}", e.getMessage());
        });
        sseEmitter.onCompletion(() -> {
            remove(userId);
            log.info("sse completed: {}", userId);
        });

        sseClients.put(userId, sseEmitter);
        log.info("new sse client: {}", userId);
        return sseEmitter;
    }

    public static void remove(String userId) {
        sseClients.remove(userId);
        log.info("remove sse client: {}", userId);
    }

    public static void send(String userId, String data, SSEMsgType type) {
        SseEmitter sseEmitter = sseClients.get(userId);
        if (sseEmitter == null) {
            log.warn("sse client not found: {}", userId);
            return;
        }
        send(sseEmitter, userId, data, type);
    }

    public static void sendAll(String data) {
        if (CollectionUtil.isEmpty(sseClients)) {
            return;
        }
        for (Map.Entry<String, SseEmitter> entry : sseClients.entrySet()) {
            String curUserId = entry.getKey();
            SseEmitter sseEmitter = entry.getValue();
            send(sseEmitter, curUserId, data, SSEMsgType.MESSAGE);
        }
    }

    private static void send(SseEmitter sseEmitter, String userId, String data, SSEMsgType type) {
        try {
            SseEmitter.SseEventBuilder event = SseEmitter.event()
                    .id(userId)
                    .data(data)
                    .name(type.name());
            sseEmitter.send(event);
        } catch (Exception e) {
            remove(userId);
            log.error("send sse error: {}", e.getMessage());
        }
    }
}
